{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ce00e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp auto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57869174",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1167bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "from os import cpu_count\n",
    "import torch\n",
    "\n",
    "from ray import tune\n",
    "from ray.tune.search.basic_variant import BasicVariantGenerator\n",
    "\n",
    "from neuralforecast.common._base_auto import BaseAuto\n",
    "from neuralforecast.common._base_auto import MockTrial\n",
    "\n",
    "from neuralforecast.models.rnn import RNN\n",
    "from neuralforecast.models.gru import GRU\n",
    "from neuralforecast.models.tcn import TCN\n",
    "from neuralforecast.models.lstm import LSTM\n",
    "from neuralforecast.models.deepar import DeepAR\n",
    "from neuralforecast.models.dilated_rnn import DilatedRNN\n",
    "from neuralforecast.models.bitcn import BiTCN\n",
    "\n",
    "from neuralforecast.models.mlp import MLP\n",
    "from neuralforecast.models.nbeats import NBEATS\n",
    "from neuralforecast.models.nbeatsx import NBEATSx\n",
    "from neuralforecast.models.nhits import NHITS\n",
    "from neuralforecast.models.dlinear import DLinear\n",
    "from neuralforecast.models.nlinear import NLinear\n",
    "from neuralforecast.models.tide import TiDE\n",
    "from neuralforecast.models.deepnpts import DeepNPTS\n",
    "\n",
    "from neuralforecast.models.tft import TFT\n",
    "from neuralforecast.models.vanillatransformer import VanillaTransformer\n",
    "from neuralforecast.models.informer import Informer\n",
    "from neuralforecast.models.autoformer import Autoformer\n",
    "from neuralforecast.models.fedformer import FEDformer\n",
    "from neuralforecast.models.patchtst import PatchTST\n",
    "from neuralforecast.models.timesnet import TimesNet\n",
    "from neuralforecast.models.itransformer import iTransformer\n",
    "\n",
    "from neuralforecast.models.stemgnn import StemGNN\n",
    "from neuralforecast.models.hint import HINT\n",
    "from neuralforecast.models.tsmixer import TSMixer\n",
    "from neuralforecast.models.tsmixerx import TSMixerx\n",
    "from neuralforecast.models.mlpmultivariate import MLPMultivariate\n",
    "from neuralforecast.models.softs import SOFTS\n",
    "\n",
    "from neuralforecast.losses.pytorch import MAE, MQLoss, DistributionLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "495dd890",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from fastcore.test import test_eq\n",
    "from nbdev.showdoc import show_doc\n",
    "\n",
    "import logging\n",
    "import warnings\n",
    "import inspect\n",
    "\n",
    "from neuralforecast.losses.pytorch import MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "490b7416",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "plt.rcParams[\"axes.grid\"]=True\n",
    "plt.rcParams['font.family'] = 'serif'\n",
    "plt.rcParams[\"figure.figsize\"] = (6,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bac3f877",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto class.\n",
    "\n",
    "# Patch for Python 3.11 on get arg spec\n",
    "if not hasattr(inspect, 'getargspec'):\n",
    "    getargspec = inspect.getfullargspec\n",
    "else:\n",
    "    getargspec = inspect.getargspec\n",
    "\n",
    "def test_args(auto_model, exclude_args=None):\n",
    "    base_auto_args = getargspec(BaseAuto)[0]\n",
    "    auto_model_args = getargspec(auto_model)[0]\n",
    "    if exclude_args is not None:\n",
    "        base_auto_args = [arg for arg in base_auto_args if arg not in exclude_args]\n",
    "    args_diff = set(base_auto_args) - set(auto_model_args)\n",
    "    assert not args_diff, f\"__init__ of {auto_model.__name__} does not contain the following required variables from BaseAuto class:\\n\\t\\t{args_diff}\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7ae65ca7",
   "metadata": {},
   "source": [
    "# Models\n",
    "\n",
    "> NeuralForecast contains user-friendly implementations of neural forecasting models that allow for easy transition of computing capabilities (GPU/CPU), computation parallelization, and hyperparameter tuning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc0437d9-7e44-4e6d-bf6b-43c2e9e338f1",
   "metadata": {},
   "source": [
    "All the NeuralForecast models are \"global\" because we train them with all the series from the input pd.DataFrame data `Y_df`, yet the optimization objective is, momentarily, \"univariate\" as it does not consider the interaction between the output predictions across time series. Like the StatsForecast library, `core.NeuralForecast` allows you to explore collections of models efficiently and contains functions for convenient wrangling of input and output pd.DataFrames predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbf7bd22",
   "metadata": {},
   "source": [
    "First we load the AirPassengers dataset such that you can run all the examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b67d291b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast.tsdataset import TimeSeriesDataset\n",
    "from neuralforecast.utils import AirPassengersDF as Y_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95293f37",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Split train/test and declare time series dataset\n",
    "Y_train_df = Y_df[Y_df.ds<='1959-12-31'] # 132 train\n",
    "Y_test_df = Y_df[Y_df.ds>'1959-12-31']   # 12 test\n",
    "dataset, *_ = TimeSeriesDataset.from_df(Y_train_df)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5b8183a9",
   "metadata": {},
   "source": [
    "# 1. Automatic Forecasting"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "858e6a1b",
   "metadata": {},
   "source": [
    "## A. RNN-Based"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60458256",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoRNN(BaseAuto):\n",
    "    \n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [-1, 4, 16, 64],\n",
    "        \"inference_input_size_multiplier\": [-1],\n",
    "        \"h\": None,\n",
    "        \"encoder_hidden_size\": tune.choice([50, 100, 200, 300]),\n",
    "        \"encoder_n_layers\": tune.randint(1, 4),\n",
    "        \"context_size\": tune.choice([5, 10, 50]),\n",
    "        \"decoder_hidden_size\": tune.choice([64, 128, 256, 512]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([16, 32]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20)\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                ):\n",
    "        \"\"\" Auto RNN\n",
    "        \n",
    "        **Parameters:**<br>\n",
    "        \n",
    "        \"\"\"\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)         \n",
    "\n",
    "        super(AutoRNN, self).__init__(\n",
    "              cls_model=RNN, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config, \n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['inference_input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['inference_input_size_multiplier']])\n",
    "        del config['input_size_multiplier'], config['inference_input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "142f7028",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoRNN, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce08be9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoRNN.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=-1, encoder_hidden_size=8)\n",
    "model = AutoRNN(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoRNN(h=12, config=None, num_samples=1, cpus=1, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ad79ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoRNN, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoRNN.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': -1, 'encoder_hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoRNN(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoRNN.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = -1\n",
    "my_config['encoder_hidden_size'] = 8\n",
    "model = AutoRNN(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "821e7999",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoLSTM(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [-1, 4, 16, 64],\n",
    "        \"inference_input_size_multiplier\": [-1],\n",
    "        \"h\": None,\n",
    "        \"encoder_hidden_size\": tune.choice([50, 100, 200, 300]),\n",
    "        \"encoder_n_layers\": tune.randint(1, 4),\n",
    "        \"context_size\": tune.choice([5, 10, 50]),\n",
    "        \"decoder_hidden_size\": tune.choice([64, 128, 256, 512]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([16, 32]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20)\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)         \n",
    "\n",
    "        super(AutoLSTM, self).__init__(\n",
    "              cls_model=LSTM,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['inference_input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['inference_input_size_multiplier']])\n",
    "        del config['input_size_multiplier'], config['inference_input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1d3965f",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoLSTM, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c68f5a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoLSTM.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=-1, encoder_hidden_size=8)\n",
    "model = AutoLSTM(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoLSTM(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71ca902",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoLSTM, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoLSTM.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': -1, 'encoder_hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoLSTM(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoLSTM.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = -1\n",
    "my_config['encoder_hidden_size'] = 8\n",
    "model = AutoLSTM(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dcaf5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoGRU(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [-1, 4, 16, 64],\n",
    "        \"inference_input_size_multiplier\": [-1],\n",
    "        \"h\": None,\n",
    "        \"encoder_hidden_size\": tune.choice([50, 100, 200, 300]),\n",
    "        \"encoder_n_layers\": tune.randint(1, 4),\n",
    "        \"context_size\": tune.choice([5, 10, 50]),\n",
    "        \"decoder_hidden_size\": tune.choice([64, 128, 256, 512]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([16, 32]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20)\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)         \n",
    "\n",
    "        super(AutoGRU, self).__init__(\n",
    "              cls_model=GRU,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config, \n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['inference_input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['inference_input_size_multiplier']])\n",
    "        del config['input_size_multiplier'], config['inference_input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b078228a",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoGRU, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf08e56",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoGRU.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=-1, encoder_hidden_size=8)\n",
    "model = AutoGRU(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoGRU(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f53b50c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoGRU, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoGRU.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': -1, 'encoder_hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoGRU(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoGRU.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = -1\n",
    "my_config['encoder_hidden_size'] = 8\n",
    "model = AutoGRU(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d5ed05",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoTCN(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [-1, 4, 16, 64],\n",
    "        \"inference_input_size_multiplier\": [-1],\n",
    "        \"h\": None,\n",
    "        \"encoder_hidden_size\": tune.choice([50, 100, 200, 300]),\n",
    "        \"context_size\": tune.choice([5, 10, 50]),\n",
    "        \"decoder_hidden_size\": tune.choice([64, 128]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([16, 32]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20)\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)         \n",
    "\n",
    "        super(AutoTCN, self).__init__(\n",
    "              cls_model=TCN,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['inference_input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['inference_input_size_multiplier']])\n",
    "        del config['input_size_multiplier'], config['inference_input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fc5d09e",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoTCN, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb7032bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoTCN.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=-1, encoder_hidden_size=8)\n",
    "model = AutoTCN(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoTCN(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cda47eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoTCN, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoTCN.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': -1, 'encoder_hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoTCN(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoTCN.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = -1\n",
    "my_config['encoder_hidden_size'] = 8\n",
    "model = AutoTCN(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b04f4d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoDeepAR(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"lstm_hidden_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"lstm_n_layers\": tune.randint(1, 4),\n",
    "        \"lstm_dropout\": tune.uniform(0.0, 0.5),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice(['robust', 'minmax1']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=DistributionLoss(distribution='StudentT', level=[80, 90], return_params=False),\n",
    "                 valid_loss=MQLoss(level=[80, 90]),\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoDeepAR, self).__init__(\n",
    "              cls_model=DeepAR, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132c8547",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoDeepAR, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61c8a07b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, lstm_hidden_size=8)\n",
    "model = AutoDeepAR(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoDeepAR(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5aec4e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoDeepAR, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoDeepAR.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'lstm_hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoDeepAR(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoDeepAR.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['lstm_hidden_size'] = 8\n",
    "model = AutoDeepAR(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b60f73",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoDilatedRNN(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [-1, 4, 16, 64],\n",
    "        \"inference_input_size_multiplier\": [-1],\n",
    "        \"h\": None,\n",
    "        \"cell_type\": tune.choice(['LSTM', 'GRU']),\n",
    "        \"encoder_hidden_size\": tune.choice([50, 100, 200, 300]),\n",
    "        \"dilations\": tune.choice([ [[1, 2], [4, 8]], [[1, 2, 4, 8]] ]),\n",
    "        \"context_size\": tune.choice([5, 10, 50]),\n",
    "        \"decoder_hidden_size\": tune.choice([64, 128, 256, 512]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([16, 32]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20)\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)         \n",
    "\n",
    "        super(AutoDilatedRNN, self).__init__(\n",
    "              cls_model=DilatedRNN,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "         )\n",
    "        \n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['inference_input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['inference_input_size_multiplier']])\n",
    "        del config['input_size_multiplier'], config['inference_input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d132351d",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoDilatedRNN, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "533eb1da",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoDilatedRNN.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=-1, encoder_hidden_size=8)\n",
    "model = AutoDilatedRNN(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoDilatedRNN(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0efe303d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoDilatedRNN, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoDilatedRNN.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': -1, 'encoder_hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoDilatedRNN(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoDilatedRNN.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = -1\n",
    "my_config['encoder_hidden_size'] = 8\n",
    "model = AutoDilatedRNN(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a0616ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoBiTCN(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([16, 32]),\n",
    "        \"dropout\": tune.uniform(0.0, 0.99),  \n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoBiTCN, self).__init__(\n",
    "              cls_model=BiTCN, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "433d2ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoBiTCN, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95850f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=8)\n",
    "model = AutoBiTCN(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoBiTCN(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c905530",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoBiTCN, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoBiTCN.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoBiTCN(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoBiTCN.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 8\n",
    "model = AutoBiTCN(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7feff99a",
   "metadata": {},
   "source": [
    "## B. MLP-Based"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa01d378",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoMLP(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice( [256, 512, 1024] ),\n",
    "        \"num_layers\": tune.randint(2, 6),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,     \n",
    "                 config=None,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "\n",
    "        # Define search space, input/output sizes       \n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoMLP, self).__init__(\n",
    "              cls_model=MLP,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config, \n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c8ecf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoMLP, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c637d57f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoMLP.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12, hidden_size=8)\n",
    "model = AutoMLP(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoMLP(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d6b5e08",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoMLP, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoMLP.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoMLP(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoMLP.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 8\n",
    "model = AutoMLP(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ed4c88b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoNBEATS(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes \n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoNBEATS, self).__init__(\n",
    "              cls_model=NBEATS, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef824625",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoNBEATS, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0f0fe9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNBEATS.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12,\n",
    "              mlp_units=3*[[8, 8]])\n",
    "model = AutoNBEATS(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoNBEATS(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7afac66d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoNBEATS, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoNBEATS.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12, 'mlp_units': 3 * [[8, 8]]})\n",
    "    return config\n",
    "\n",
    "model = AutoNBEATS(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoNBEATS.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['mlp_units'] = 3 * [[8, 8]]\n",
    "model = AutoNBEATS(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "753c5300",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoNBEATSx(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoNBEATSx, self).__init__(\n",
    "              cls_model=NBEATSx,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2829115",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoNBEATSx, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca94b63",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNBEATS.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12,\n",
    "              mlp_units=3*[[8, 8]])\n",
    "model = AutoNBEATSx(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoNBEATSx(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76334984",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoNBEATSx, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoNBEATSx.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12, 'mlp_units': 3 * [[8, 8]]})\n",
    "    return config\n",
    "\n",
    "model = AutoNBEATSx(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoNBEATSx.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['mlp_units'] = 3 * [[8, 8]]\n",
    "model = AutoNBEATSx(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79a2ede0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoNHITS(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "       \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "       \"h\": None,\n",
    "       \"n_pool_kernel_size\": tune.choice([[2, 2, 1], 3*[1], 3*[2], 3*[4], \n",
    "                                         [8, 4, 1], [16, 8, 1]]),\n",
    "       \"n_freq_downsample\": tune.choice([[168, 24, 1], [24, 12, 1], \n",
    "                                         [180, 60, 1], [60, 8, 1], \n",
    "                                         [40, 20, 1], [1, 1, 1]]),\n",
    "       \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "       \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "       \"max_steps\": tune.quniform(lower=500, upper=1500, q=100),\n",
    "       \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "       \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "       \"loss\": None,\n",
    "       \"random_seed\": tune.randint(lower=1, upper=20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                ):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)                     \n",
    "\n",
    "        super(AutoNHITS, self).__init__(\n",
    "              cls_model=NHITS, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd1236fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoNHITS, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc67276b-6005-4f67-969d-36a90130e0f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12, \n",
    "              mlp_units=3 * [[8, 8]])\n",
    "model = AutoNHITS(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoNHITS(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa9c4e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoNHITS, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoNHITS.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12, 'mlp_units': 3 * [[8, 8]]})\n",
    "    return config\n",
    "\n",
    "model = AutoNHITS(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoNHITS.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['mlp_units'] = 3 * [[8, 8]]\n",
    "model = AutoNHITS(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b5616d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoDLinear(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "       \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "       \"h\": None,\n",
    "       \"moving_avg_window\": tune.choice([11, 25, 51]),\n",
    "       \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "       \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "       \"max_steps\": tune.quniform(lower=500, upper=1500, q=100),\n",
    "       \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "       \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "       \"loss\": None,\n",
    "       \"random_seed\": tune.randint(lower=1, upper=20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                ):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)                   \n",
    "\n",
    "        super(AutoDLinear, self).__init__(\n",
    "              cls_model=DLinear, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04a1aab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoDLinear, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc8b0d4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoDLinear.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12)\n",
    "model = AutoDLinear(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoDLinear(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcf93b8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoDLinear, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoDLinear.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoDLinear(h=12, config=my_config_new, backend='optuna', cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoDLinear.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoDLinear(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "180a5a10",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoNLinear(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "       \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "       \"h\": None,\n",
    "       \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "       \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "       \"max_steps\": tune.quniform(lower=500, upper=1500, q=100),\n",
    "       \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "       \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "       \"loss\": None,\n",
    "       \"random_seed\": tune.randint(lower=1, upper=20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                ):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)                 \n",
    "\n",
    "        super(AutoNLinear, self).__init__(\n",
    "              cls_model=NLinear, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e414458",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoNLinear, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80759148",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNLinear.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12)\n",
    "model = AutoNLinear(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoNLinear(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a36677a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoNLinear, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoNLinear.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoNLinear(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoNLinear.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoNLinear(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "973a470e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoTiDE(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "       \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "       \"h\": None,\n",
    "       \"hidden_size\": tune.choice([256, 512, 1024]),\n",
    "       \"decoder_output_dim\": tune.choice([8, 16, 32]),\n",
    "       \"temporal_decoder_dim\": tune.choice([32, 64, 128]),\n",
    "       \"num_encoder_layers\": tune.choice([1, 2, 3]),\n",
    "       \"num_decoder_layers\": tune.choice([1, 2, 3]),\n",
    "       \"temporal_width\": tune.choice([4, 8, 16]),\n",
    "       \"dropout\":tune.choice([0.0, 0.1, 0.2, 0.3, 0.5]),\n",
    "       \"layernorm\": tune.choice([True, False]),\n",
    "       \"learning_rate\": tune.loguniform(1e-5, 1e-2),\n",
    "       \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "       \"max_steps\": tune.quniform(lower=500, upper=1500, q=100),\n",
    "       \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "       \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "       \"loss\": None,\n",
    "       \"random_seed\": tune.randint(lower=1, upper=20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                ):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)                 \n",
    "\n",
    "        super(AutoTiDE, self).__init__(\n",
    "              cls_model=TiDE, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d31d3bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoTiDE, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ae8f192",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoTiDE.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12)\n",
    "model = AutoTiDE(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoTiDE(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d66600b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoTiDE, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoTiDE.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoTiDE(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoTiDE.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoTiDE(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "890e455f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoDeepNPTS(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "       \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "       \"h\": None,\n",
    "       \"hidden_size\": tune.choice([16, 32, 64]),\n",
    "       \"dropout\": tune.uniform(0.0, 0.99),\n",
    "       \"n_layers\": tune.choice([1, 2, 4]),\n",
    "       \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "       \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "       \"max_steps\": tune.quniform(lower=500, upper=1500, q=100),\n",
    "       \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "       \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "       \"loss\": None,\n",
    "       \"random_seed\": tune.randint(lower=1, upper=20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                ):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)                 \n",
    "\n",
    "        super(AutoDeepNPTS, self).__init__(\n",
    "              cls_model=DeepNPTS, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples,\n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afa62ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoDeepNPTS, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1b6c380",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoDeepNPTS.default_config\n",
    "config = dict(max_steps=2, val_check_steps=1, input_size=12)\n",
    "model = AutoDeepNPTS(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoDeepNPTS(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68f5d045",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoDeepNPTS, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoDeepNPTS.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoDeepNPTS(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoDeepNPTS.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoDeepNPTS(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fd705a56",
   "metadata": {},
   "source": [
    "## C. Transformer-Based"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91cd9b23",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoTFT(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256]),\n",
    "        \"n_head\": tune.choice([4, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)             \n",
    "\n",
    "        super(AutoTFT, self).__init__(\n",
    "              cls_model=TFT, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ab1077",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoTFT, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e283e96f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=8)\n",
    "model = AutoTFT(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoTFT(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a78492a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoTFT, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoTFT.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoTFT(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoTFT.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 8\n",
    "model = AutoTFT(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28b5ea6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoVanillaTransformer(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256]),\n",
    "        \"n_head\": tune.choice([4, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoVanillaTransformer, self).__init__(\n",
    "              cls_model=VanillaTransformer, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "265a8197",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoVanillaTransformer, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3f3e429",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=8)\n",
    "model = AutoVanillaTransformer(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoVanillaTransformer(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb6a1ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoVanillaTransformer, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoVanillaTransformer.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoVanillaTransformer(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoVanillaTransformer.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 8\n",
    "model = AutoVanillaTransformer(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9d509ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoInformer(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256]),\n",
    "        \"n_head\": tune.choice([4, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoInformer, self).__init__(\n",
    "              cls_model=Informer, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c43869",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoInformer, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1648c77b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=8)\n",
    "model = AutoInformer(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoInformer(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c70e2e8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoInformer, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoInformer.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoInformer(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoInformer.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 8\n",
    "model = AutoInformer(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d409d246",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoAutoformer(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256]),\n",
    "        \"n_head\": tune.choice([4, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoAutoformer, self).__init__(\n",
    "              cls_model=Autoformer, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e34860",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoAutoformer, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38d4b216",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=8)\n",
    "model = AutoAutoformer(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoAutoformer(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee9e62a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoAutoformer, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoAutoformer.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 8})\n",
    "    return config\n",
    "\n",
    "model = AutoAutoformer(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoAutoformer.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 8\n",
    "model = AutoAutoformer(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c84a19",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoFEDformer(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes    \n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)          \n",
    "\n",
    "        super(AutoFEDformer, self).__init__(\n",
    "              cls_model=FEDformer, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45497479",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoFEDformer, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b7b1ccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=64)\n",
    "model = AutoFEDformer(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoFEDformer(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8b8a02d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoFEDformer, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoFEDformer.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 64})\n",
    "    return config\n",
    "\n",
    "model = AutoFEDformer(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoFEDformer.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 64\n",
    "model = AutoFEDformer(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb3cde6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoPatchTST(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([16, 128, 256]),\n",
    "        \"n_heads\": tune.choice([4, 16]),\n",
    "        \"patch_len\": tune.choice([16, 24]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"revin\": tune.choice([False, True]),\n",
    "        \"max_steps\": tune.choice([500, 1000, 5000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"windows_batch_size\": tune.choice([128, 256, 512, 1024]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)           \n",
    "\n",
    "        super(AutoPatchTST, self).__init__(\n",
    "              cls_model=PatchTST, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config['input_size_multiplier']])  \n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d807209",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoPatchTST, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13878573",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoNHITS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=16)\n",
    "model = AutoPatchTST(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoPatchTST(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eff3267c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoPatchTST, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoPatchTST.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 16})\n",
    "    return config\n",
    "\n",
    "model = AutoPatchTST(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoPatchTST.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 16\n",
    "model = AutoPatchTST(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61c3be9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoiTransformer(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"n_series\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256]),\n",
    "        \"n_heads\": tune.choice([4, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 n_series,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)                 \n",
    "\n",
    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
    "        if backend == 'ray':\n",
    "            config['n_series'] = n_series\n",
    "        elif backend == 'optuna':\n",
    "            mock_trial = MockTrial()\n",
    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
    "                raise Exception(f\"config needs 'n_series': {n_series}\")           \n",
    "\n",
    "        super(AutoiTransformer, self).__init__(\n",
    "              cls_model=iTransformer, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series):\n",
    "        config = cls.default_config.copy()        \n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config[\"input_size_multiplier\"]])\n",
    "\n",
    "        # Rolling windows with step_size=1 or step_size=h\n",
    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
    "        config['step_size'] = tune.choice([1, h])\n",
    "        del config[\"input_size_multiplier\"]\n",
    "        if backend == 'optuna':\n",
    "            # Always use n_series from parameters\n",
    "            config['n_series'] = n_series\n",
    "            config = cls._ray_config_to_optuna(config)           \n",
    "\n",
    "        return config         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f416fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoiTransformer, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffd40db",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoiTransformer.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=16)\n",
    "model = AutoiTransformer(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoiTransformer(h=12, n_series=1, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a2052de",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoiTransformer, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoiTransformer.get_default_config(h=12, n_series=1, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 16})\n",
    "    return config\n",
    "\n",
    "model = AutoiTransformer(h=12, n_series=1, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoiTransformer.get_default_config(h=12, n_series=1, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 16\n",
    "model = AutoiTransformer(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "57d6cb1f",
   "metadata": {},
   "source": [
    "## D. CNN Based"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f775426",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoTimesNet(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"hidden_size\": tune.choice([32, 64, 128]),\n",
    "        \"conv_hidden_size\": tune.choice([32, 64, 128]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice(['robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128]),\n",
    "        \"windows_batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend)                  \n",
    "\n",
    "        super(AutoTimesNet, self).__init__(\n",
    "              cls_model=TimesNet, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        config = cls.default_config.copy()\n",
    "        config['input_size'] = tune.choice([h*x \\\n",
    "                        for x in config['input_size_multiplier']])\n",
    "        config['step_size'] = tune.choice([1, h])        \n",
    "        del config['input_size_multiplier']\n",
    "        if backend == 'optuna':\n",
    "            config = cls._ray_config_to_optuna(config)         \n",
    "\n",
    "        return config         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2312754",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoTimesNet, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1467a471",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoTimesNet.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=32)\n",
    "model = AutoTimesNet(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoTimesNet(h=12, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b60eb58",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoTimesNet, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoTimesNet.get_default_config(h=12, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 2, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 32})\n",
    "    return config\n",
    "\n",
    "model = AutoTimesNet(h=12, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoTimesNet.get_default_config(h=12, backend='ray')\n",
    "my_config['max_steps'] = 2\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 32\n",
    "model = AutoTimesNet(h=12, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e6fd22c7",
   "metadata": {},
   "source": [
    "## E. Multivariate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6784e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoStemGNN(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4],\n",
    "        \"h\": None,\n",
    "        \"n_series\": None,\n",
    "        \"n_stacks\": tune.choice([2]),\n",
    "        \"multi_layer\": tune.choice([3, 5, 7]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 n_series,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)                  \n",
    "\n",
    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
    "        if backend == 'ray':\n",
    "            config['n_series'] = n_series\n",
    "        elif backend == 'optuna':\n",
    "            mock_trial = MockTrial()\n",
    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
    "                raise Exception(f\"config needs 'n_series': {n_series}\")\n",
    "\n",
    "        super(AutoStemGNN, self).__init__(\n",
    "              cls_model=StemGNN, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series):\n",
    "        config = cls.default_config.copy()        \n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config[\"input_size_multiplier\"]])\n",
    "\n",
    "        # Rolling windows with step_size=1 or step_size=h\n",
    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
    "        config['step_size'] = tune.choice([1, h])\n",
    "        del config[\"input_size_multiplier\"]\n",
    "        if backend == 'optuna':\n",
    "            # Always use n_series from parameters\n",
    "            config['n_series'] = n_series\n",
    "            config = cls._ray_config_to_optuna(config)           \n",
    "\n",
    "        return config        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1163c1de",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoStemGNN, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca99c2ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoStemGNN.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12)\n",
    "model = AutoStemGNN(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoStemGNN(h=12, n_series=1, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cd03851",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "assert model.config(MockTrial())['n_series'] == 1\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoStemGNN, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoStemGNN.get_default_config(h=12, backend='optuna', n_series=1)\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoStemGNN(h=12, n_series=1, config=my_config_new, backend='optuna')\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoStemGNN.get_default_config(h=12, backend='ray', n_series=1)\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoStemGNN(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58844aca",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoHINT(BaseAuto):\n",
    "\n",
    "    def __init__(self,\n",
    "                 cls_model,\n",
    "                 h,\n",
    "                 loss,\n",
    "                 valid_loss,\n",
    "                 S,\n",
    "                 config,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 refit_with_val=False,\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None,\n",
    "                 ):\n",
    "        \n",
    "        super(AutoHINT, self).__init__(\n",
    "              cls_model=cls_model, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,\n",
    "        )\n",
    "        if backend == 'optuna':\n",
    "            raise Exception(\"Optuna is not supported for AutoHINT.\")\n",
    "\n",
    "        # Validate presence of reconciliation strategy\n",
    "        # parameter in configuration space\n",
    "        if not ('reconciliation' in config.keys()):\n",
    "            raise Exception(\"config needs reconciliation, \\\n",
    "                            try tune.choice(['BottomUp', 'MinTraceOLS', 'MinTraceWLS'])\")\n",
    "        self.S = S\n",
    "\n",
    "    def _fit_model(self, cls_model, config,\n",
    "                   dataset, val_size, test_size, distributed_config=None):\n",
    "        # Overwrite _fit_model for HINT two-stage instantiation\n",
    "        reconciliation = config.pop('reconciliation')\n",
    "        base_model = cls_model(**config)\n",
    "        model = HINT(h=base_model.h, model=base_model, \n",
    "                     S=self.S, reconciliation=reconciliation)\n",
    "        model.test_size = test_size\n",
    "        model = model.fit(\n",
    "            dataset,\n",
    "            val_size=val_size, \n",
    "            test_size=test_size,\n",
    "            distributed_config=distributed_config,\n",
    "        )\n",
    "        return model\n",
    "    \n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series=None):\n",
    "        raise Exception(\"AutoHINT has no default configuration.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "622a0888",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoHINT, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf210740",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "def sort_df_hier(Y_df, S_df):\n",
    "    # NeuralForecast core, sorts unique_id lexicographically\n",
    "    # by default, this class matches S_df and Y_hat_df order.    \n",
    "    Y_df.unique_id = Y_df.unique_id.astype('category')\n",
    "    Y_df.unique_id = Y_df.unique_id.cat.set_categories(S_df.index)\n",
    "    Y_df = Y_df.sort_values(by=['unique_id', 'ds'])\n",
    "    return Y_df\n",
    "\n",
    "# -----Create synthetic dataset-----\n",
    "np.random.seed(123)\n",
    "train_steps = 20\n",
    "num_levels = 7\n",
    "level = np.arange(0, 100, 0.1)\n",
    "qs = [[50-lv/2, 50+lv/2] for lv in level]\n",
    "quantiles = np.sort(np.concatenate(qs)/100)\n",
    "\n",
    "levels = ['Top', 'Mid1', 'Mid2', 'Bottom1', 'Bottom2', 'Bottom3', 'Bottom4']\n",
    "unique_ids = np.repeat(levels, train_steps)\n",
    "\n",
    "S = np.array([[1., 1., 1., 1.],\n",
    "              [1., 1., 0., 0.],\n",
    "              [0., 0., 1., 1.],\n",
    "              [1., 0., 0., 0.],\n",
    "              [0., 1., 0., 0.],\n",
    "              [0., 0., 1., 0.],\n",
    "              [0., 0., 0., 1.]])\n",
    "\n",
    "S_dict = {col: S[:, i] for i, col in enumerate(levels[3:])}\n",
    "S_df = pd.DataFrame(S_dict, index=levels)\n",
    "\n",
    "ds = pd.date_range(start='2018-03-31', periods=train_steps, freq='Q').tolist() * num_levels\n",
    "# Create Y_df\n",
    "y_lists = [S @ np.random.uniform(low=100, high=500, size=4) for i in range(train_steps)]\n",
    "y = [elem for tup in zip(*y_lists) for elem in tup]\n",
    "Y_df = pd.DataFrame({'unique_id': unique_ids, 'ds': ds, 'y': y})\n",
    "Y_df = sort_df_hier(Y_df, S_df)\n",
    "\n",
    "hint_dataset, *_ = TimeSeriesDataset.from_df(df=Y_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d41f07c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Perform a simple hyperparameter optimization with \n",
    "# NHITS and then reconcile with HINT\n",
    "from neuralforecast.losses.pytorch import GMM, sCRPS\n",
    "\n",
    "base_config = dict(max_steps=1, val_check_steps=1, input_size=8)\n",
    "base_model = AutoNHITS(h=4, loss=GMM(n_components=2, quantiles=quantiles), \n",
    "                       config=base_config, num_samples=1, cpus=1)\n",
    "model = HINT(h=4, S=S_df.values,\n",
    "             model=base_model,  reconciliation='MinTraceOLS')\n",
    "\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=hint_dataset)\n",
    "\n",
    "# Perform a conjunct hyperparameter optimization with \n",
    "# NHITS + HINT reconciliation configurations\n",
    "nhits_config = {\n",
    "       \"learning_rate\": tune.choice([1e-3]),                                     # Initial Learning rate\n",
    "       \"max_steps\": tune.choice([1]),                                            # Number of SGD steps\n",
    "       \"val_check_steps\": tune.choice([1]),                                      # Number of steps between validation\n",
    "       \"input_size\": tune.choice([5 * 12]),                                      # input_size = multiplier * horizon\n",
    "       \"batch_size\": tune.choice([7]),                                           # Number of series in windows\n",
    "       \"windows_batch_size\": tune.choice([256]),                                 # Number of windows in batch\n",
    "       \"n_pool_kernel_size\": tune.choice([[2, 2, 2], [16, 8, 1]]),               # MaxPool's Kernelsize\n",
    "       \"n_freq_downsample\": tune.choice([[168, 24, 1], [24, 12, 1], [1, 1, 1]]), # Interpolation expressivity ratios\n",
    "       \"activation\": tune.choice(['ReLU']),                                      # Type of non-linear activation\n",
    "       \"n_blocks\":  tune.choice([[1, 1, 1]]),                                    # Blocks per each 3 stacks\n",
    "       \"mlp_units\":  tune.choice([[[512, 512], [512, 512], [512, 512]]]),        # 2 512-Layers per block for each stack\n",
    "       \"interpolation_mode\": tune.choice(['linear']),                            # Type of multi-step interpolation\n",
    "       \"random_seed\": tune.randint(1, 10),\n",
    "       \"reconciliation\": tune.choice(['BottomUp', 'MinTraceOLS', 'MinTraceWLS'])\n",
    "    }\n",
    "model = AutoHINT(h=4, S=S_df.values,\n",
    "                 cls_model=NHITS,\n",
    "                 config=nhits_config,\n",
    "                 loss=GMM(n_components=2, level=[80, 90]),\n",
    "                 valid_loss=sCRPS(level=[80, 90]),\n",
    "                 num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=hint_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58183a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoHINT) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfe91bb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoTSMixer(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4],\n",
    "        \"h\": None,\n",
    "        \"n_series\": None,\n",
    "        \"n_block\": tune.choice([1, 2, 4, 6, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-2),\n",
    "        \"ff_dim\": tune.choice([32, 64, 128]),\n",
    "        \"scaler_type\": tune.choice(['identity', 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"dropout\": tune.uniform(0.0, 0.99),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 n_series,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)                 \n",
    "\n",
    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
    "        if backend == 'ray':\n",
    "            config['n_series'] = n_series\n",
    "        elif backend == 'optuna':\n",
    "            mock_trial = MockTrial()\n",
    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
    "                raise Exception(f\"config needs 'n_series': {n_series}\")\n",
    "\n",
    "        super(AutoTSMixer, self).__init__(\n",
    "              cls_model=TSMixer, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series):\n",
    "        config = cls.default_config.copy()        \n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config[\"input_size_multiplier\"]])\n",
    "\n",
    "        # Rolling windows with step_size=1 or step_size=h\n",
    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
    "        config['step_size'] = tune.choice([1, h])\n",
    "        del config[\"input_size_multiplier\"]\n",
    "        if backend == 'optuna':\n",
    "            # Always use n_series from parameters\n",
    "            config['n_series'] = n_series\n",
    "            config = cls._ray_config_to_optuna(config)           \n",
    "\n",
    "        return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b844b1be",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoTSMixer, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3cafb38",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoTSMixer.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12)\n",
    "model = AutoTSMixer(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoTSMixer(h=12, n_series=1, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "062abe64",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "assert model.config(MockTrial())['n_series'] == 1\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoTSMixer, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoTSMixer.get_default_config(h=12, backend='optuna', n_series=1)\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoTSMixer(h=12, n_series=1, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoTSMixer.get_default_config(h=12, backend='ray', n_series=1)\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoTSMixer(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a83fec63",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoTSMixerx(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4],\n",
    "        \"h\": None,\n",
    "        \"n_series\": None,\n",
    "        \"n_block\": tune.choice([1, 2, 4, 6, 8]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-2),\n",
    "        \"ff_dim\": tune.choice([32, 64, 128]),\n",
    "        \"scaler_type\": tune.choice(['identity', 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"dropout\": tune.uniform(0.0, 0.99),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 n_series,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)         \n",
    "\n",
    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
    "        if backend == 'ray':\n",
    "            config['n_series'] = n_series\n",
    "        elif backend == 'optuna':\n",
    "            mock_trial = MockTrial()\n",
    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
    "                raise Exception(f\"config needs 'n_series': {n_series}\")            \n",
    "\n",
    "        super(AutoTSMixerx, self).__init__(\n",
    "              cls_model=TSMixerx, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series):\n",
    "        config = cls.default_config.copy()        \n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config[\"input_size_multiplier\"]])\n",
    "\n",
    "        # Rolling windows with step_size=1 or step_size=h\n",
    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
    "        config['step_size'] = tune.choice([1, h])\n",
    "        del config[\"input_size_multiplier\"]\n",
    "        if backend == 'optuna':\n",
    "            # Always use n_series from parameters\n",
    "            config['n_series'] = n_series\n",
    "            config = cls._ray_config_to_optuna(config)           \n",
    "\n",
    "        return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f99bd4dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoTSMixerx, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41f67b64",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoTSMixerx.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12)\n",
    "model = AutoTSMixerx(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoTSMixerx(h=12, n_series=1, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05210f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "assert model.config(MockTrial())['n_series'] == 1\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoTSMixerx, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoTSMixerx.get_default_config(h=12, backend='optuna', n_series=1)\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoTSMixerx(h=12, n_series=1, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoTSMixerx.get_default_config(h=12, backend='ray', n_series=1)\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoTSMixerx(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18518b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoMLPMultivariate(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"n_series\": None,\n",
    "        \"hidden_size\": tune.choice( [256, 512, 1024] ),\n",
    "        \"num_layers\": tune.randint(2, 6),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard']),\n",
    "        \"max_steps\": tune.choice([500, 1000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 n_series,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,     \n",
    "                 config=None,\n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "\n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)         \n",
    "\n",
    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
    "        if backend == 'ray':\n",
    "            config['n_series'] = n_series\n",
    "        elif backend == 'optuna':\n",
    "            mock_trial = MockTrial()\n",
    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
    "                raise Exception(f\"config needs 'n_series': {n_series}\")  \n",
    "\n",
    "        super(AutoMLPMultivariate, self).__init__(\n",
    "              cls_model=MLPMultivariate,\n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config, \n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series):\n",
    "        config = cls.default_config.copy()        \n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config[\"input_size_multiplier\"]])\n",
    "\n",
    "        # Rolling windows with step_size=1 or step_size=h\n",
    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
    "        config['step_size'] = tune.choice([1, h])\n",
    "        del config[\"input_size_multiplier\"]\n",
    "        if backend == 'optuna':\n",
    "            # Always use n_series from parameters\n",
    "            config['n_series'] = n_series\n",
    "            config = cls._ray_config_to_optuna(config)           \n",
    "\n",
    "        return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d077b8e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(AutoMLPMultivariate, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd40c5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoTSMixerx.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12)\n",
    "model = AutoMLPMultivariate(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoMLPMultivariate(h=12, n_series=1, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64639b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "assert model.config(MockTrial())['n_series'] == 1\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoMLPMultivariate, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoMLPMultivariate.get_default_config(h=12, backend='optuna', n_series=1)\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12})\n",
    "    return config\n",
    "\n",
    "model = AutoMLPMultivariate(h=12, n_series=1, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoMLPMultivariate.get_default_config(h=12, backend='ray', n_series=1)\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "model = AutoMLPMultivariate(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33495c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class AutoSOFTS(BaseAuto):\n",
    "\n",
    "    default_config = {\n",
    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
    "        \"h\": None,\n",
    "        \"n_series\": None,\n",
    "        \"hidden_size\": tune.choice([64, 128, 256, 512]),\n",
    "        \"d_core\": tune.choice([64, 128, 256, 512]),\n",
    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
    "        \"scaler_type\": tune.choice([None, 'robust', 'standard', 'identity']),\n",
    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
    "        \"loss\": None,\n",
    "        \"random_seed\": tune.randint(1, 20),\n",
    "    }\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 n_series,\n",
    "                 loss=MAE(),\n",
    "                 valid_loss=None,\n",
    "                 config=None, \n",
    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
    "                 num_samples=10,\n",
    "                 refit_with_val=False,\n",
    "                 cpus=cpu_count(),\n",
    "                 gpus=torch.cuda.device_count(),\n",
    "                 verbose=False,\n",
    "                 alias=None,\n",
    "                 backend='ray',\n",
    "                 callbacks=None):\n",
    "        \n",
    "        # Define search space, input/output sizes\n",
    "        if config is None:\n",
    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)                 \n",
    "\n",
    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
    "        if backend == 'ray':\n",
    "            config['n_series'] = n_series\n",
    "        elif backend == 'optuna':\n",
    "            mock_trial = MockTrial()\n",
    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
    "                raise Exception(f\"config needs 'n_series': {n_series}\")           \n",
    "\n",
    "        super(AutoSOFTS, self).__init__(\n",
    "              cls_model=SOFTS, \n",
    "              h=h,\n",
    "              loss=loss,\n",
    "              valid_loss=valid_loss,\n",
    "              config=config,\n",
    "              search_alg=search_alg,\n",
    "              num_samples=num_samples, \n",
    "              refit_with_val=refit_with_val,\n",
    "              cpus=cpus,\n",
    "              gpus=gpus,\n",
    "              verbose=verbose,\n",
    "              alias=alias,\n",
    "              backend=backend,\n",
    "              callbacks=callbacks,            \n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def get_default_config(cls, h, backend, n_series):\n",
    "        config = cls.default_config.copy()        \n",
    "        config['input_size'] = tune.choice([h * x \\\n",
    "                        for x in config[\"input_size_multiplier\"]])\n",
    "\n",
    "        # Rolling windows with step_size=1 or step_size=h\n",
    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
    "        config['step_size'] = tune.choice([1, h])\n",
    "        del config[\"input_size_multiplier\"]\n",
    "        if backend == 'optuna':\n",
    "            # Always use n_series from parameters\n",
    "            config['n_series'] = n_series\n",
    "            config = cls._ray_config_to_optuna(config)           \n",
    "\n",
    "        return config         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd2bc7bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/Nixtla/neuralforecast/blob/main/neuralforecast/auto.py#L2085){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### AutoSOFTS\n",
       "\n",
       ">      AutoSOFTS (h, n_series, loss=MAE(), valid_loss=None, config=None,\n",
       ">                 search_alg=<ray.tune.search.basic_variant.BasicVariantGenerato\n",
       ">                 r object at 0x32ce6d180>, num_samples=10,\n",
       ">                 refit_with_val=False, cpus=11, gpus=0, verbose=False,\n",
       ">                 alias=None, backend='ray', callbacks=None)\n",
       "\n",
       "Class for Automatic Hyperparameter Optimization, it builds on top of `ray` to\n",
       "give access to a wide variety of hyperparameter optimization tools ranging\n",
       "from classic grid search, to Bayesian optimization and HyperBand algorithm.\n",
       "\n",
       "The validation loss to be optimized is defined by the `config['loss']` dictionary\n",
       "value, the config also contains the rest of the hyperparameter search space.\n",
       "\n",
       "It is important to note that the success of this hyperparameter optimization\n",
       "heavily relies on a strong correlation between the validation and test periods.\n",
       "\n",
       "|    | **Type** | **Default** | **Details** |\n",
       "| -- | -------- | ----------- | ----------- |\n",
       "| h | int |  | Forecast horizon |\n",
       "| n_series |  |  |  |\n",
       "| loss | MAE | MAE() | Instantiated train loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html). |\n",
       "| valid_loss | NoneType | None | Instantiated valid loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html). |\n",
       "| config | NoneType | None | Dictionary with ray.tune defined search space or function that takes an optuna trial and returns a configuration dict. |\n",
       "| search_alg | BasicVariantGenerator | <ray.tune.search.basic_variant.BasicVariantGenerator object> | For ray see https://docs.ray.io/en/latest/tune/api_docs/suggestion.html<br>For optuna see https://optuna.readthedocs.io/en/stable/reference/samplers/index.html. |\n",
       "| num_samples | int | 10 | Number of hyperparameter optimization steps/samples. |\n",
       "| refit_with_val | bool | False | Refit of best model should preserve val_size. |\n",
       "| cpus | int | 11 | Number of cpus to use during optimization. Only used with ray tune. |\n",
       "| gpus | int | 0 | Number of gpus to use during optimization, default all available. Only used with ray tune. |\n",
       "| verbose | bool | False | Track progress. |\n",
       "| alias | NoneType | None | Custom name of the model. |\n",
       "| backend | str | ray | Backend to use for searching the hyperparameter space, can be either 'ray' or 'optuna'. |\n",
       "| callbacks | NoneType | None | List of functions to call during the optimization process.<br>ray reference: https://docs.ray.io/en/latest/tune/tutorials/tune-metrics.html<br>optuna reference: https://optuna.readthedocs.io/en/stable/tutorial/20_recipes/007_optuna_callback.html |"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/Nixtla/neuralforecast/blob/main/neuralforecast/auto.py#L2085){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### AutoSOFTS\n",
       "\n",
       ">      AutoSOFTS (h, n_series, loss=MAE(), valid_loss=None, config=None,\n",
       ">                 search_alg=<ray.tune.search.basic_variant.BasicVariantGenerato\n",
       ">                 r object at 0x32ce6d180>, num_samples=10,\n",
       ">                 refit_with_val=False, cpus=11, gpus=0, verbose=False,\n",
       ">                 alias=None, backend='ray', callbacks=None)\n",
       "\n",
       "Class for Automatic Hyperparameter Optimization, it builds on top of `ray` to\n",
       "give access to a wide variety of hyperparameter optimization tools ranging\n",
       "from classic grid search, to Bayesian optimization and HyperBand algorithm.\n",
       "\n",
       "The validation loss to be optimized is defined by the `config['loss']` dictionary\n",
       "value, the config also contains the rest of the hyperparameter search space.\n",
       "\n",
       "It is important to note that the success of this hyperparameter optimization\n",
       "heavily relies on a strong correlation between the validation and test periods.\n",
       "\n",
       "|    | **Type** | **Default** | **Details** |\n",
       "| -- | -------- | ----------- | ----------- |\n",
       "| h | int |  | Forecast horizon |\n",
       "| n_series |  |  |  |\n",
       "| loss | MAE | MAE() | Instantiated train loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html). |\n",
       "| valid_loss | NoneType | None | Instantiated valid loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html). |\n",
       "| config | NoneType | None | Dictionary with ray.tune defined search space or function that takes an optuna trial and returns a configuration dict. |\n",
       "| search_alg | BasicVariantGenerator | <ray.tune.search.basic_variant.BasicVariantGenerator object> | For ray see https://docs.ray.io/en/latest/tune/api_docs/suggestion.html<br>For optuna see https://optuna.readthedocs.io/en/stable/reference/samplers/index.html. |\n",
       "| num_samples | int | 10 | Number of hyperparameter optimization steps/samples. |\n",
       "| refit_with_val | bool | False | Refit of best model should preserve val_size. |\n",
       "| cpus | int | 11 | Number of cpus to use during optimization. Only used with ray tune. |\n",
       "| gpus | int | 0 | Number of gpus to use during optimization, default all available. Only used with ray tune. |\n",
       "| verbose | bool | False | Track progress. |\n",
       "| alias | NoneType | None | Custom name of the model. |\n",
       "| backend | str | ray | Backend to use for searching the hyperparameter space, can be either 'ray' or 'optuna'. |\n",
       "| callbacks | NoneType | None | List of functions to call during the optimization process.<br>ray reference: https://docs.ray.io/en/latest/tune/tutorials/tune-metrics.html<br>optuna reference: https://optuna.readthedocs.io/en/stable/tutorial/20_recipes/007_optuna_callback.html |"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(AutoSOFTS, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43050009",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-31 15:06:51,623\tINFO worker.py:1749 -- Started a local Ray instance.\n",
      "2024-05-31 15:06:52,217\tINFO tune.py:263 -- Initializing Ray automatically. For cluster usage or custom Ray initialization, call `ray.init(...)` before `Tuner(...)`.\n",
      "2024-05-31 15:06:52,219\tINFO tune.py:624 -- [output] This uses the legacy output and progress reporter, as Jupyter notebooks are not supported by the new engine, yet. For more information, please see https://github.com/ray-project/ray/issues/36949\n",
      "2024-05-31 15:06:54,633\tERROR tune_controller.py:1332 -- Trial task failed for trial _train_tune_eb814_00000\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/air/execution/_internal/event_manager.py\", line 110, in resolve_future\n",
      "    result = ray.get(future)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/_private/auto_init_hook.py\", line 21, in auto_init_wrapper\n",
      "    return fn(*args, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/_private/client_mode_hook.py\", line 103, in wrapper\n",
      "    return func(*args, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/_private/worker.py\", line 2623, in get\n",
      "    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/_private/worker.py\", line 861, in get_objects\n",
      "    raise value.as_instanceof_cause()\n",
      "ray.exceptions.RayTaskError(ValueError): \u001b[36mray::ImplicitFunc.train()\u001b[39m (pid=3951, ip=127.0.0.1, actor_id=1641e0248328a9ca364b098901000000, repr=_train_tune)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/tune/trainable/trainable.py\", line 334, in train\n",
      "    raise skipped from exception_cause(skipped)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/air/_internal/util.py\", line 98, in run\n",
      "    self._ret = self._target(*self._args, **self._kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/tune/trainable/function_trainable.py\", line 53, in <lambda>\n",
      "    training_func=lambda: self._trainable_func(self.config),\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/tune/trainable/function_trainable.py\", line 261, in _trainable_func\n",
      "    output = fn()\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/ray/tune/trainable/util.py\", line 130, in inner\n",
      "    return trainable(config, **fn_kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py\", line 209, in _train_tune\n",
      "    _ = self._fit_model(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py\", line 357, in _fit_model\n",
      "    model = model.fit(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py\", line 541, in fit\n",
      "    return self._fit(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_model.py\", line 281, in _fit\n",
      "    trainer = pl.Trainer(**model.trainer_kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/utilities/argparse.py\", line 70, in insert_env_defaults\n",
      "    return fn(self, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py\", line 431, in __init__\n",
      "    self._callback_connector.on_trainer_init(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/callback_connector.py\", line 79, in on_trainer_init\n",
      "    _validate_callbacks_list(self.trainer.callbacks)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/callback_connector.py\", line 227, in _validate_callbacks_list\n",
      "    stateful_callbacks = [cb for cb in callbacks if is_overridden(\"state_dict\", instance=cb)]\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/callback_connector.py\", line 227, in <listcomp>\n",
      "    stateful_callbacks = [cb for cb in callbacks if is_overridden(\"state_dict\", instance=cb)]\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/utilities/model_helpers.py\", line 39, in is_overridden\n",
      "    raise ValueError(\"Expected a parent\")\n",
      "ValueError: Expected a parent\n",
      "2024-05-31 15:06:54,638\tWARNING experiment_state.py:205 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n",
      "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n",
      "You can suppress this error by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0).\n",
      "2024-05-31 15:06:54,640\tINFO tune.py:1021 -- Wrote the latest version of all result files and experiment state to '/Users/marcopeix/ray_results/_train_tune_2024-05-31_15-06-49' in 0.0031s.\n",
      "2024-05-31 15:06:54,641\tERROR tune.py:1049 -- Trials did not complete: [_train_tune_eb814_00000]\n",
      "Seed set to 1\n",
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "`Trainer(val_check_interval=1)` was configured so validation will run after every batch.\n",
      "`Trainer.fit` stopped: `max_steps=1` reached.\n",
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "`Trainer(val_check_interval=1)` was configured so validation will run after every batch.\n"
     ]
    }
   ],
   "source": [
    "%%capture\n",
    "# Use your own config or AutoSOFTS.default_config\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12, hidden_size=16)\n",
    "model = AutoSOFTS(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Optuna\n",
    "model = AutoSOFTS(h=12, n_series=1, config=None, backend='optuna')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2588d91b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-05-31 15:07:11,775] A new study created in memory with name: no-name-9172b813-168d-48bf-9ae8-6d9dc8603e2b\n",
      "Seed set to 17\n",
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "`Trainer(val_check_interval=1)` was configured so validation will run after every batch.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16737e68ad76462e9854de99d232b32d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[W 2024-05-31 15:07:11,926] Trial 0 failed with parameters: {'hidden_size': 256, 'd_core': 256, 'learning_rate': 0.0003957532292586766, 'scaler_type': 'robust', 'max_steps': 2000, 'batch_size': 32, 'random_seed': 17, 'input_size': 24, 'step_size': 1} because of the following error: NotImplementedError(\"The operator 'aten::nanmedian.dim_values' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.\").\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/optuna/study/_optimize.py\", line 196, in _run_trial\n",
      "    value_or_values = func(trial)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py\", line 320, in objective\n",
      "    model = self._fit_model(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py\", line 357, in _fit_model\n",
      "    model = model.fit(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py\", line 541, in fit\n",
      "    return self._fit(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_model.py\", line 282, in _fit\n",
      "    trainer.fit(model, datamodule=datamodule)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py\", line 544, in fit\n",
      "    call._call_and_handle_interrupt(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py\", line 44, in _call_and_handle_interrupt\n",
      "    return trainer_fn(*args, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py\", line 580, in _fit_impl\n",
      "    self._run(model, ckpt_path=ckpt_path)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py\", line 989, in _run\n",
      "    results = self._run_stage()\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py\", line 1033, in _run_stage\n",
      "    self._run_sanity_check()\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py\", line 1062, in _run_sanity_check\n",
      "    val_loop.run()\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py\", line 182, in _decorator\n",
      "    return loop_run(self, *args, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py\", line 134, in run\n",
      "    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py\", line 391, in _evaluation_step\n",
      "    output = call._call_strategy_hook(trainer, hook_name, *step_args)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py\", line 309, in _call_strategy_hook\n",
      "    output = fn(*args, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py\", line 403, in validation_step\n",
      "    return self.lightning_module.validation_step(*args, **kwargs)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py\", line 407, in validation_step\n",
      "    windows = self._normalization(windows=windows, y_idx=y_idx)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py\", line 264, in _normalization\n",
      "    temporal_data = self.scaler.transform(x=temporal_data, mask=temporal_mask)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_scalers.py\", line 424, in transform\n",
      "    x_shift, x_scale = self.compute_statistics(\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_scalers.py\", line 224, in robust_statistics\n",
      "    x_median = masked_median(x=x, mask=mask, dim=dim)\n",
      "  File \"/Users/marcopeix/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_scalers.py\", line 30, in masked_median\n",
      "    x_median, _ = x_nan.nanmedian(dim=dim, keepdim=keepdim)\n",
      "NotImplementedError: The operator 'aten::nanmedian.dim_values' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.\n",
      "[W 2024-05-31 15:07:11,928] Trial 0 failed with value None.\n"
     ]
    },
    {
     "ename": "NotImplementedError",
     "evalue": "The operator 'aten::nanmedian.dim_values' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 16\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m config\n\u001b[1;32m     15\u001b[0m model \u001b[38;5;241m=\u001b[39m AutoSOFTS(h\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m12\u001b[39m, n_series\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, config\u001b[38;5;241m=\u001b[39mmy_config_new, backend\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124moptuna\u001b[39m\u001b[38;5;124m'\u001b[39m, num_samples\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, cpus\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 16\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdataset\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdataset\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# Unit test for situation: Ray with updated default config\u001b[39;00m\n\u001b[1;32m     19\u001b[0m my_config \u001b[38;5;241m=\u001b[39m AutoSOFTS\u001b[38;5;241m.\u001b[39mget_default_config(h\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m12\u001b[39m, n_series\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, backend\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mray\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py:412\u001b[0m, in \u001b[0;36mBaseAuto.fit\u001b[0;34m(self, dataset, val_size, test_size, random_seed, distributed_config)\u001b[0m\n\u001b[1;32m    410\u001b[0m     best_config \u001b[38;5;241m=\u001b[39m results\u001b[38;5;241m.\u001b[39mget_best_result()\u001b[38;5;241m.\u001b[39mconfig\n\u001b[1;32m    411\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 412\u001b[0m     results \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_optuna_tune_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    413\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcls_model\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcls_model\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    414\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdataset\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    415\u001b[0m \u001b[43m        \u001b[49m\u001b[43mval_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mval_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    416\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtest_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtest_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    417\u001b[0m \u001b[43m        \u001b[49m\u001b[43mverbose\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mverbose\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    418\u001b[0m \u001b[43m        \u001b[49m\u001b[43mnum_samples\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnum_samples\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    419\u001b[0m \u001b[43m        \u001b[49m\u001b[43msearch_alg\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msearch_alg\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    420\u001b[0m \u001b[43m        \u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    421\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdistributed_config\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdistributed_config\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    422\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    423\u001b[0m     best_config \u001b[38;5;241m=\u001b[39m results\u001b[38;5;241m.\u001b[39mbest_trial\u001b[38;5;241m.\u001b[39muser_attrs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mALL_PARAMS\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m    424\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_fit_model(\n\u001b[1;32m    425\u001b[0m     cls_model\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcls_model,\n\u001b[1;32m    426\u001b[0m     config\u001b[38;5;241m=\u001b[39mbest_config,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    430\u001b[0m     distributed_config\u001b[38;5;241m=\u001b[39mdistributed_config,\n\u001b[1;32m    431\u001b[0m )\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py:345\u001b[0m, in \u001b[0;36mBaseAuto._optuna_tune_model\u001b[0;34m(self, cls_model, dataset, val_size, test_size, verbose, num_samples, search_alg, config, distributed_config)\u001b[0m\n\u001b[1;32m    342\u001b[0m     sampler \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    344\u001b[0m study \u001b[38;5;241m=\u001b[39m optuna\u001b[38;5;241m.\u001b[39mcreate_study(sampler\u001b[38;5;241m=\u001b[39msampler, direction\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mminimize\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 345\u001b[0m \u001b[43mstudy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptimize\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    346\u001b[0m \u001b[43m    \u001b[49m\u001b[43mobjective\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    347\u001b[0m \u001b[43m    \u001b[49m\u001b[43mn_trials\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_samples\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mshow_progress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mverbose\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcallbacks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    350\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    351\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m study\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/optuna/study/study.py:451\u001b[0m, in \u001b[0;36mStudy.optimize\u001b[0;34m(self, func, n_trials, timeout, n_jobs, catch, callbacks, gc_after_trial, show_progress_bar)\u001b[0m\n\u001b[1;32m    348\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21moptimize\u001b[39m(\n\u001b[1;32m    349\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    350\u001b[0m     func: ObjectiveFuncType,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    357\u001b[0m     show_progress_bar: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    358\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    359\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Optimize an objective function.\u001b[39;00m\n\u001b[1;32m    360\u001b[0m \n\u001b[1;32m    361\u001b[0m \u001b[38;5;124;03m    Optimization is done by choosing a suitable set of hyperparameter values from a given\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    449\u001b[0m \u001b[38;5;124;03m            If nested invocation of this method occurs.\u001b[39;00m\n\u001b[1;32m    450\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 451\u001b[0m     \u001b[43m_optimize\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    452\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstudy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    453\u001b[0m \u001b[43m        \u001b[49m\u001b[43mfunc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    454\u001b[0m \u001b[43m        \u001b[49m\u001b[43mn_trials\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_trials\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    455\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    456\u001b[0m \u001b[43m        \u001b[49m\u001b[43mn_jobs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_jobs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    457\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcatch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mtuple\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mcatch\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43misinstance\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mcatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mIterable\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43mcatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    458\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallbacks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    459\u001b[0m \u001b[43m        \u001b[49m\u001b[43mgc_after_trial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgc_after_trial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    460\u001b[0m \u001b[43m        \u001b[49m\u001b[43mshow_progress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mshow_progress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    461\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/optuna/study/_optimize.py:62\u001b[0m, in \u001b[0;36m_optimize\u001b[0;34m(study, func, n_trials, timeout, n_jobs, catch, callbacks, gc_after_trial, show_progress_bar)\u001b[0m\n\u001b[1;32m     60\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m     61\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m n_jobs \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m---> 62\u001b[0m         \u001b[43m_optimize_sequential\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     63\u001b[0m \u001b[43m            \u001b[49m\u001b[43mstudy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     64\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     65\u001b[0m \u001b[43m            \u001b[49m\u001b[43mn_trials\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[43m            \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     67\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcatch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     68\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     69\u001b[0m \u001b[43m            \u001b[49m\u001b[43mgc_after_trial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     70\u001b[0m \u001b[43m            \u001b[49m\u001b[43mreseed_sampler_rng\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     71\u001b[0m \u001b[43m            \u001b[49m\u001b[43mtime_start\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     72\u001b[0m \u001b[43m            \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     75\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m n_jobs \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m:\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/optuna/study/_optimize.py:159\u001b[0m, in \u001b[0;36m_optimize_sequential\u001b[0;34m(study, func, n_trials, timeout, catch, callbacks, gc_after_trial, reseed_sampler_rng, time_start, progress_bar)\u001b[0m\n\u001b[1;32m    156\u001b[0m         \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m    158\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 159\u001b[0m     frozen_trial \u001b[38;5;241m=\u001b[39m \u001b[43m_run_trial\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstudy\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    160\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    161\u001b[0m     \u001b[38;5;66;03m# The following line mitigates memory problems that can be occurred in some\u001b[39;00m\n\u001b[1;32m    162\u001b[0m     \u001b[38;5;66;03m# environments (e.g., services that use computing containers such as GitHub Actions).\u001b[39;00m\n\u001b[1;32m    163\u001b[0m     \u001b[38;5;66;03m# Please refer to the following PR for further details:\u001b[39;00m\n\u001b[1;32m    164\u001b[0m     \u001b[38;5;66;03m# https://github.com/optuna/optuna/pull/325.\u001b[39;00m\n\u001b[1;32m    165\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m gc_after_trial:\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/optuna/study/_optimize.py:247\u001b[0m, in \u001b[0;36m_run_trial\u001b[0;34m(study, func, catch)\u001b[0m\n\u001b[1;32m    240\u001b[0m         \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28;01mFalse\u001b[39;00m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mShould not reach.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    242\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    243\u001b[0m     frozen_trial\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;241m==\u001b[39m TrialState\u001b[38;5;241m.\u001b[39mFAIL\n\u001b[1;32m    244\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m func_err \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    245\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(func_err, catch)\n\u001b[1;32m    246\u001b[0m ):\n\u001b[0;32m--> 247\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m func_err\n\u001b[1;32m    248\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m frozen_trial\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/optuna/study/_optimize.py:196\u001b[0m, in \u001b[0;36m_run_trial\u001b[0;34m(study, func, catch)\u001b[0m\n\u001b[1;32m    194\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m get_heartbeat_thread(trial\u001b[38;5;241m.\u001b[39m_trial_id, study\u001b[38;5;241m.\u001b[39m_storage):\n\u001b[1;32m    195\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 196\u001b[0m         value_or_values \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    197\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m exceptions\u001b[38;5;241m.\u001b[39mTrialPruned \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    198\u001b[0m         \u001b[38;5;66;03m# TODO(mamu): Handle multi-objective cases.\u001b[39;00m\n\u001b[1;32m    199\u001b[0m         state \u001b[38;5;241m=\u001b[39m TrialState\u001b[38;5;241m.\u001b[39mPRUNED\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py:320\u001b[0m, in \u001b[0;36mBaseAuto._optuna_tune_model.<locals>.objective\u001b[0;34m(trial)\u001b[0m\n\u001b[1;32m    318\u001b[0m user_cfg \u001b[38;5;241m=\u001b[39m config(trial)\n\u001b[1;32m    319\u001b[0m cfg \u001b[38;5;241m=\u001b[39m deepcopy(user_cfg)\n\u001b[0;32m--> 320\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_fit_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    321\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcls_model\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcls_model\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    322\u001b[0m \u001b[43m    \u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcfg\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    323\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdataset\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    324\u001b[0m \u001b[43m    \u001b[49m\u001b[43mval_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mval_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    325\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtest_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtest_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    326\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdistributed_config\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdistributed_config\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    327\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    328\u001b[0m trial\u001b[38;5;241m.\u001b[39mset_user_attr(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mALL_PARAMS\u001b[39m\u001b[38;5;124m\"\u001b[39m, user_cfg)\n\u001b[1;32m    329\u001b[0m metrics \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mmetrics\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_auto.py:357\u001b[0m, in \u001b[0;36mBaseAuto._fit_model\u001b[0;34m(self, cls_model, config, dataset, val_size, test_size, distributed_config)\u001b[0m\n\u001b[1;32m    353\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_fit_model\u001b[39m(\n\u001b[1;32m    354\u001b[0m     \u001b[38;5;28mself\u001b[39m, cls_model, config, dataset, val_size, test_size, distributed_config\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    355\u001b[0m ):\n\u001b[1;32m    356\u001b[0m     model \u001b[38;5;241m=\u001b[39m cls_model(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mconfig)\n\u001b[0;32m--> 357\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    358\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    359\u001b[0m \u001b[43m        \u001b[49m\u001b[43mval_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mval_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    360\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtest_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtest_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    361\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdistributed_config\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdistributed_config\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    362\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    363\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m model\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py:541\u001b[0m, in \u001b[0;36mBaseMultivariate.fit\u001b[0;34m(self, dataset, val_size, test_size, random_seed, distributed_config)\u001b[0m\n\u001b[1;32m    537\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m distributed_config \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    538\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    539\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmultivariate models cannot be trained using distributed data parallel.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    540\u001b[0m     )\n\u001b[0;32m--> 541\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_fit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    542\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdataset\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    543\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_series\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    544\u001b[0m \u001b[43m    \u001b[49m\u001b[43mvalid_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_series\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    545\u001b[0m \u001b[43m    \u001b[49m\u001b[43mval_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mval_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    546\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtest_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtest_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    547\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrandom_seed\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrandom_seed\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    548\u001b[0m \u001b[43m    \u001b[49m\u001b[43mshuffle_train\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    549\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdistributed_config\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    550\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_model.py:282\u001b[0m, in \u001b[0;36mBaseModel._fit\u001b[0;34m(self, dataset, batch_size, valid_batch_size, val_size, test_size, random_seed, shuffle_train, distributed_config)\u001b[0m\n\u001b[1;32m    280\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m    281\u001b[0m trainer \u001b[38;5;241m=\u001b[39m pl\u001b[38;5;241m.\u001b[39mTrainer(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mmodel\u001b[38;5;241m.\u001b[39mtrainer_kwargs)\n\u001b[0;32m--> 282\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdatamodule\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdatamodule\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    283\u001b[0m model\u001b[38;5;241m.\u001b[39mmetrics \u001b[38;5;241m=\u001b[39m trainer\u001b[38;5;241m.\u001b[39mcallback_metrics\n\u001b[1;32m    284\u001b[0m model\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__dict__\u001b[39m\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_trainer\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:544\u001b[0m, in \u001b[0;36mTrainer.fit\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    542\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mstatus \u001b[38;5;241m=\u001b[39m TrainerStatus\u001b[38;5;241m.\u001b[39mRUNNING\n\u001b[1;32m    543\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m--> 544\u001b[0m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_and_handle_interrupt\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    545\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_fit_impl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_dataloaders\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_dataloaders\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdatamodule\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mckpt_path\u001b[49m\n\u001b[1;32m    546\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:44\u001b[0m, in \u001b[0;36m_call_and_handle_interrupt\u001b[0;34m(trainer, trainer_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m     42\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mlauncher \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     43\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mlauncher\u001b[38;5;241m.\u001b[39mlaunch(trainer_fn, \u001b[38;5;241m*\u001b[39margs, trainer\u001b[38;5;241m=\u001b[39mtrainer, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 44\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrainer_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     46\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _TunerExitException:\n\u001b[1;32m     47\u001b[0m     _call_teardown_hook(trainer)\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:580\u001b[0m, in \u001b[0;36mTrainer._fit_impl\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    573\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mfn \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    574\u001b[0m ckpt_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_checkpoint_connector\u001b[38;5;241m.\u001b[39m_select_ckpt_path(\n\u001b[1;32m    575\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mfn,\n\u001b[1;32m    576\u001b[0m     ckpt_path,\n\u001b[1;32m    577\u001b[0m     model_provided\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    578\u001b[0m     model_connected\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    579\u001b[0m )\n\u001b[0;32m--> 580\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mckpt_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mckpt_path\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    582\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mstopped\n\u001b[1;32m    583\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:989\u001b[0m, in \u001b[0;36mTrainer._run\u001b[0;34m(self, model, ckpt_path)\u001b[0m\n\u001b[1;32m    984\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_signal_connector\u001b[38;5;241m.\u001b[39mregister_signal_handlers()\n\u001b[1;32m    986\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[1;32m    987\u001b[0m \u001b[38;5;66;03m# RUN THE TRAINER\u001b[39;00m\n\u001b[1;32m    988\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[0;32m--> 989\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_stage\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    991\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[1;32m    992\u001b[0m \u001b[38;5;66;03m# POST-Training CLEAN UP\u001b[39;00m\n\u001b[1;32m    993\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[1;32m    994\u001b[0m log\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m: trainer tearing down\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1033\u001b[0m, in \u001b[0;36mTrainer._run_stage\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1031\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining:\n\u001b[1;32m   1032\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m isolate_rng():\n\u001b[0;32m-> 1033\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_sanity_check\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1034\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mautograd\u001b[38;5;241m.\u001b[39mset_detect_anomaly(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_detect_anomaly):\n\u001b[1;32m   1035\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfit_loop\u001b[38;5;241m.\u001b[39mrun()\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1062\u001b[0m, in \u001b[0;36mTrainer._run_sanity_check\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1059\u001b[0m call\u001b[38;5;241m.\u001b[39m_call_callback_hooks(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mon_sanity_check_start\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   1061\u001b[0m \u001b[38;5;66;03m# run eval step\u001b[39;00m\n\u001b[0;32m-> 1062\u001b[0m \u001b[43mval_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1064\u001b[0m call\u001b[38;5;241m.\u001b[39m_call_callback_hooks(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mon_sanity_check_end\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   1066\u001b[0m \u001b[38;5;66;03m# reset logger connector\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py:182\u001b[0m, in \u001b[0;36m_no_grad_context.<locals>._decorator\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    180\u001b[0m     context_manager \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mno_grad\n\u001b[1;32m    181\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m context_manager():\n\u001b[0;32m--> 182\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mloop_run\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:134\u001b[0m, in \u001b[0;36m_EvaluationLoop.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    132\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbatch_progress\u001b[38;5;241m.\u001b[39mis_last_batch \u001b[38;5;241m=\u001b[39m data_fetcher\u001b[38;5;241m.\u001b[39mdone\n\u001b[1;32m    133\u001b[0m     \u001b[38;5;66;03m# run step hooks\u001b[39;00m\n\u001b[0;32m--> 134\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_evaluation_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataloader_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataloader_iter\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    135\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m:\n\u001b[1;32m    136\u001b[0m     \u001b[38;5;66;03m# this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support\u001b[39;00m\n\u001b[1;32m    137\u001b[0m     \u001b[38;5;28;01mbreak\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:391\u001b[0m, in \u001b[0;36m_EvaluationLoop._evaluation_step\u001b[0;34m(self, batch, batch_idx, dataloader_idx, dataloader_iter)\u001b[0m\n\u001b[1;32m    385\u001b[0m hook_name \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtest_step\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mtesting \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvalidation_step\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    386\u001b[0m step_args \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    387\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_build_step_args_from_hook_kwargs(hook_kwargs, hook_name)\n\u001b[1;32m    388\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m using_dataloader_iter\n\u001b[1;32m    389\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m (dataloader_iter,)\n\u001b[1;32m    390\u001b[0m )\n\u001b[0;32m--> 391\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_strategy_hook\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrainer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhook_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mstep_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    393\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbatch_progress\u001b[38;5;241m.\u001b[39mincrement_processed()\n\u001b[1;32m    395\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m using_dataloader_iter:\n\u001b[1;32m    396\u001b[0m     \u001b[38;5;66;03m# update the hook kwargs now that the step method might have consumed the iterator\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:309\u001b[0m, in \u001b[0;36m_call_strategy_hook\u001b[0;34m(trainer, hook_name, *args, **kwargs)\u001b[0m\n\u001b[1;32m    306\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    308\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[Strategy]\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtrainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mhook_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 309\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    311\u001b[0m \u001b[38;5;66;03m# restore current_fx when nested context\u001b[39;00m\n\u001b[1;32m    312\u001b[0m pl_module\u001b[38;5;241m.\u001b[39m_current_fx_name \u001b[38;5;241m=\u001b[39m prev_fx_name\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py:403\u001b[0m, in \u001b[0;36mStrategy.validation_step\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    401\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module:\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_redirection(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvalidation_step\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 403\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlightning_module\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalidation_step\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py:407\u001b[0m, in \u001b[0;36mBaseMultivariate.validation_step\u001b[0;34m(self, batch, batch_idx)\u001b[0m\n\u001b[1;32m    405\u001b[0m windows \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_create_windows(batch, step\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mval\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    406\u001b[0m y_idx \u001b[38;5;241m=\u001b[39m batch[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124my_idx\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[0;32m--> 407\u001b[0m windows \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_normalization\u001b[49m\u001b[43m(\u001b[49m\u001b[43mwindows\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwindows\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_idx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43my_idx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    409\u001b[0m \u001b[38;5;66;03m# Parse windows\u001b[39;00m\n\u001b[1;32m    410\u001b[0m (\n\u001b[1;32m    411\u001b[0m     insample_y,\n\u001b[1;32m    412\u001b[0m     insample_mask,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    417\u001b[0m     stat_exog,\n\u001b[1;32m    418\u001b[0m ) \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_parse_windows(batch, windows)\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_base_multivariate.py:264\u001b[0m, in \u001b[0;36mBaseMultivariate._normalization\u001b[0;34m(self, windows, y_idx)\u001b[0m\n\u001b[1;32m    260\u001b[0m \u001b[38;5;66;03m# Normalize. self.scaler stores the shift and scale for inverse transform\u001b[39;00m\n\u001b[1;32m    261\u001b[0m temporal_mask \u001b[38;5;241m=\u001b[39m temporal_mask\u001b[38;5;241m.\u001b[39munsqueeze(\n\u001b[1;32m    262\u001b[0m     \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    263\u001b[0m )  \u001b[38;5;66;03m# Add channel dimension for scaler.transform.\u001b[39;00m\n\u001b[0;32m--> 264\u001b[0m temporal_data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscaler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtransform\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtemporal_data\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtemporal_mask\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# Replace values in windows dict\u001b[39;00m\n\u001b[1;32m    266\u001b[0m temporal[:, temporal_idxs, :, :] \u001b[38;5;241m=\u001b[39m temporal_data\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_scalers.py:424\u001b[0m, in \u001b[0;36mTemporalNorm.transform\u001b[0;34m(self, x, mask)\u001b[0m\n\u001b[1;32m    412\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtransform\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, mask):\n\u001b[1;32m    413\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Center and scale the data.\u001b[39;00m\n\u001b[1;32m    414\u001b[0m \n\u001b[1;32m    415\u001b[0m \u001b[38;5;124;03m    **Parameters:**<br>\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    422\u001b[0m \u001b[38;5;124;03m    `z`: torch.Tensor same shape as `x`, except scaled.\u001b[39;00m\n\u001b[1;32m    423\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 424\u001b[0m     x_shift, x_scale \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompute_statistics\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    425\u001b[0m \u001b[43m        \u001b[49m\u001b[43mx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmask\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdim\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meps\u001b[49m\n\u001b[1;32m    426\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    427\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mx_shift \u001b[38;5;241m=\u001b[39m x_shift\n\u001b[1;32m    428\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mx_scale \u001b[38;5;241m=\u001b[39m x_scale\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_scalers.py:224\u001b[0m, in \u001b[0;36mrobust_statistics\u001b[0;34m(x, mask, dim, eps)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrobust_statistics\u001b[39m(x, mask, dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, eps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-6\u001b[39m):\n\u001b[1;32m    199\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Robust Median Scaler\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[1;32m    201\u001b[0m \u001b[38;5;124;03m    Standardizes features by removing the median and scaling\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    222\u001b[0m \u001b[38;5;124;03m    `z`: torch.Tensor same shape as `x`, except scaled.\u001b[39;00m\n\u001b[1;32m    223\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m     x_median \u001b[38;5;241m=\u001b[39m \u001b[43mmasked_median\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmask\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdim\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    225\u001b[0m     x_mad \u001b[38;5;241m=\u001b[39m masked_median(x\u001b[38;5;241m=\u001b[39mtorch\u001b[38;5;241m.\u001b[39mabs(x \u001b[38;5;241m-\u001b[39m x_median), mask\u001b[38;5;241m=\u001b[39mmask, dim\u001b[38;5;241m=\u001b[39mdim)\n\u001b[1;32m    227\u001b[0m     \u001b[38;5;66;03m# Protect x_mad=0 values\u001b[39;00m\n\u001b[1;32m    228\u001b[0m     \u001b[38;5;66;03m# Assuming normality and relationship between mad and std\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/neuralforecast/lib/python3.10/site-packages/neuralforecast/common/_scalers.py:30\u001b[0m, in \u001b[0;36mmasked_median\u001b[0;34m(x, mask, dim, keepdim)\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Masked Median\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \n\u001b[1;32m     15\u001b[0m \u001b[38;5;124;03mCompute the median of tensor `x` along dim, ignoring values where\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     27\u001b[0m \u001b[38;5;124;03m`x_median`: torch.Tensor with normalized values.\u001b[39;00m\n\u001b[1;32m     28\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     29\u001b[0m x_nan \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mfloat()\u001b[38;5;241m.\u001b[39mmasked_fill(mask \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m1\u001b[39m, \u001b[38;5;28mfloat\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnan\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m---> 30\u001b[0m x_median, _ \u001b[38;5;241m=\u001b[39m \u001b[43mx_nan\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnanmedian\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdim\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkeepdim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeepdim\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m x_median \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mnan_to_num(x_median, nan\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.0\u001b[39m)\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m x_median\n",
      "\u001b[0;31mNotImplementedError\u001b[0m: The operator 'aten::nanmedian.dim_values' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS."
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "# Check Optuna\n",
    "assert model.config(MockTrial())['h'] == 12\n",
    "\n",
    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
    "test_args(AutoSOFTS, exclude_args=['cls_model']) \n",
    "\n",
    "# Unit test for situation: Optuna with updated default config\n",
    "my_config = AutoSOFTS.get_default_config(h=12, n_series=1, backend='optuna')\n",
    "def my_config_new(trial):\n",
    "    config = {**my_config(trial)}\n",
    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'hidden_size': 16})\n",
    "    return config\n",
    "\n",
    "model = AutoSOFTS(h=12, n_series=1, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)\n",
    "\n",
    "# Unit test for situation: Ray with updated default config\n",
    "my_config = AutoSOFTS.get_default_config(h=12, n_series=1, backend='ray')\n",
    "my_config['max_steps'] = 1\n",
    "my_config['val_check_steps'] = 1\n",
    "my_config['input_size'] = 12\n",
    "my_config['hidden_size'] = 16\n",
    "model = AutoSOFTS(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
    "model.fit(dataset=dataset)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "69d588d9",
   "metadata": {},
   "source": [
    "# TESTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db128f64-b311-479e-bf81-07c3bf9f1a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast.tsdataset import TimeSeriesDataset\n",
    "from neuralforecast.utils import AirPassengersDF as Y_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1691fdcc-99e3-472b-ae26-03fb89847227",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Split train/test and declare time series dataset\n",
    "Y_train_df = Y_df[Y_df.ds<='1959-12-31'] # 132 train\n",
    "Y_test_df = Y_df[Y_df.ds>'1959-12-31']   # 12 test\n",
    "dataset, *_ = TimeSeriesDataset.from_df(Y_train_df)\n",
    "\n",
    "config = dict(max_steps=1, val_check_steps=1, input_size=12)\n",
    "model = AutoNHITS(h=12, config=config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "535cf3a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "## TESTS\n",
    "nhits_config = {\n",
    "       \"learning_rate\": tune.choice([1e-3]),                                     # Initial Learning rate\n",
    "       \"max_steps\": tune.choice([1]),                                            # Number of SGD steps\n",
    "       \"val_check_steps\": tune.choice([1]),                                      # Number of steps between validation\n",
    "       \"input_size\": tune.choice([5 * 12]),                                      # input_size = multiplier * horizon\n",
    "       \"batch_size\": tune.choice([7]),                                           # Number of series in windows\n",
    "       \"windows_batch_size\": tune.choice([256]),                                 # Number of windows in batch\n",
    "       \"n_pool_kernel_size\": tune.choice([[2, 2, 2], [16, 8, 1]]),               # MaxPool's Kernelsize\n",
    "       \"n_freq_downsample\": tune.choice([[168, 24, 1], [24, 12, 1], [1, 1, 1]]), # Interpolation expressivity ratios\n",
    "       \"activation\": tune.choice(['ReLU']),                                      # Type of non-linear activation\n",
    "       \"n_blocks\":  tune.choice([[1, 1, 1]]),                                    # Blocks per each 3 stacks\n",
    "       \"mlp_units\":  tune.choice([[[512, 512], [512, 512], [512, 512]]]),        # 2 512-Layers per block for each stack\n",
    "       \"interpolation_mode\": tune.choice(['linear']),                            # Type of multi-step interpolation\n",
    "       \"random_seed\": tune.randint(1, 10),\n",
    "    }\n",
    "\n",
    "model = AutoNHITS(h=12, loss=MAE(), valid_loss=MSE(), config=nhits_config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "# Test equality\n",
    "test_eq(str(type(model.valid_loss)), \"<class 'neuralforecast.losses.pytorch.MSE'>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bd1d011-aafb-4b7c-a0a0-0f190b529fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from neuralforecast.losses.pytorch import GMM, sCRPS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7e471ca-a633-4868-be0b-3f2992b6cdc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "## TODO: Add unit tests for interactions between loss/valid_loss types\n",
    "## TODO: Unit tests (2 types of networks x 2 types of loss x 2 types of valid loss)\n",
    "## Checking if base recurrent methods run point valid_loss correctly\n",
    "tcn_config = {\n",
    "       \"learning_rate\": tune.choice([1e-3]),                                     # Initial Learning rate\n",
    "       \"max_steps\": tune.choice([1]),                                            # Number of SGD steps\n",
    "       \"val_check_steps\": tune.choice([1]),                                      # Number of steps between validation\n",
    "       \"input_size\": tune.choice([5 * 12]),                                      # input_size = multiplier * horizon\n",
    "       \"batch_size\": tune.choice([7]),                                           # Number of series in windows\n",
    "       \"random_seed\": tune.randint(1, 10),\n",
    "    }\n",
    "\n",
    "model = AutoTCN(h=12, \n",
    "                loss=MAE(), \n",
    "                valid_loss=MSE(), \n",
    "                config=tcn_config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)\n",
    "\n",
    "## Checking if base recurrent methods run quantile valid_loss correctly\n",
    "model = AutoTCN(h=12, \n",
    "                loss=GMM(n_components=2, level=[80, 90]),\n",
    "                valid_loss=sCRPS(level=[80, 90]),\n",
    "                config=tcn_config, num_samples=1, cpus=1)\n",
    "\n",
    "# Fit and predict\n",
    "model.fit(dataset=dataset)\n",
    "y_hat = model.predict(dataset=dataset)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
